import sqlite3
from Parameters import *
from PolymorphismTables import PolymorphismTables
from R2_Tools import R2
from SiteFrequencySpectrum import SiteFrequencySpectrum
import infiniteSitesTools as ist 
from Polyspect_PowerSphere import PolyspectPowerSphere
from Databases import DB_Polyspect

import math
#import matplotlib.pyplot as plt
import numpy.matlib as np

import sys

def main(argv):
	
	print "***************************************************"
	print "********************* New run *********************"
	print "***************************************************"
	print ""
	
	nChr = 30 # int(argv[0])
#	segSites = 20 #int(argv[1])
#	alpha = 100.0 # float(argv[2])

	ptNull = None
	mySFS = None

	alphaToCheck = [ 100.0, 200.0, 500.0 ] #, 2.0, 5.0, 10.0, 20.0, 50.0, 100.0]
	tauToCheck = [ 200, 500, 1000 ]
	segSitesToCheck = [ x for x in range(20, 201, 20) ]
	
	for segSites in segSitesToCheck:
		for alpha in alphaToCheck:
			 for tau in tauToCheck:
	
				pop = getDefaultPopulationParameters( nChr, segSites )
				
				N1 = pop.N * math.exp(-alpha * tau / (4.0*pop.N))
				print "N1\t", N1
				
				nullDemography = Demography(pop.N)
				nullDemography.addEpoch( alpha, 0, False)
				nullDemography.addEpoch( N1, tau, True )
	
	
				print "Ho\t", nullDemography
				print pop.getMSString_Null(), nullDemography.getMSString()
	
				simParams = SimulationParameters( pop, nullDemography )
				print simParams
							
				altDemography = Demography( pop.N )
				altDemography.addEpoch( pop.N, 0, True )
				print "Ha\t", altDemography
			
				ptNull = PolymorphismTables()
				msCommand = simParams.getMSString_Null()
				msOut = ptNull.runMS(msCommand)
				ptNull.readMS(None, msOut)
		
				mySFSs = SiteFrequencySpectrum(ptNull)
				
				pt = PolymorphismTables()
				msCommand = simParams.getMSString_Alternative(altDemography )
				msOut = pt.runMS(msCommand)
				pt.readMS(None, msOut)
				myAltSFSs = SiteFrequencySpectrum(pt)
				
				db = DB_Polyspect()
				polyspects = db.getAllPolyspects( nChr )
				
				mySpace = PolyspectPowerSphere( simParams, altDemography, mySFSs, myAltSFSs, polyspects)
				
				point = mySpace.getCurrentPoint()
				print point
				
				for i in range(0,10000):
					newPoint = mySpace.getNextPoint()
					if newPoint == None:
						print "\nFound peak for now"
						break
						
					print newPoint
			
				db.addPowerPoint( mySpace )
					
#	rr = R2( simParams )
#	r2power = rr.getPower(altDemography)
#	print 'R2 power:\t' + str(r2power)
#
#	makePlot( mySpace, r2power )


def getDefaultPopulationParameters( nChr=28, segSites=50 ):
	
	popParameters = Population()
	
	popParameters.numChromosomes = nChr
	popParameters.numSegregatingSites = segSites

	popParameters.rateMutation = 1E-8
	popParameters.rateRecombination = 0.0
	popParameters.sequenceLength = 25000

	popParameters.N = 10000
	popParameters.numSimulations = 10000
	popParameters.numPowerSimulations = 1250

	return popParameters


def makePlot( mySpace, r2power ):
	tjd = ist.normalizeToFirst( mySpace.myPowerPath[0].myPolyspect.myOperator )
	flf = ist.normalizeToFirst( mySpace.myPowerPath[3].myPolyspect.myOperator )
	dmd = ist.normalizeToFirst( mySpace.myPowerPath[-1].myPolyspect.myOperator )


	myPowers = []
	myPowers.append( mySpace.myPowerPath[0].myPower)
	myPowers.append( mySpace.myPowerPath[2].myPower)
	myPowers.append( mySpace.myPowerPath[3].myPower)
	myPowers.append( r2power )
	myPowers.append( mySpace.myPowerPath[-1].myPower)
	
	fig = plt.figure()
	ind = np.arange(30)
	
	ax = fig.add_subplot(2,2,1)
	ax.bar( ind, tjd, color='m' )
	ax.set_title('Tajimas D')

	ax = fig.add_subplot(2,2,2)
	ax.bar( ind, flf, color='b')
	ax.set_title('Fu Lis F')

	ax = fig.add_subplot(2,2,3)
	ax.bar( ind, dmd, color='y')
	ax.set_title('Davids M')

	ind = np.arange(5)
	ax = fig.add_subplot(2,2,4)
	rects = ax.bar( ind, myPowers, color='r')
	ax.set_title('Powers')
	ax.set_yticks([0, 0.2, 0.4, 0.6, 0.8, 1.0])
	ax.set_xticks(ind+0.35)
	ax.set_xticklabels( ('TajD', 'FuLiD', 'FuLiF', 'R2', 'DavidsM') )

	def autolabel(rects):
		# attach some text labels
		for rect in rects:
			height = rect.get_height()
			ax.text(rect.get_x()+rect.get_width()/2., 1.05*height, '%d'%int(100*height), ha='center', va='bottom')

	autolabel(rects)

	print "\n*** Final path ***"
	print mySpace
	print "\nProgram complete"
	
	plt.show()


	#r = codebox( "*** Final path ***", "Results", str(mySpace) )
	
if __name__ == "__main__":
	main(sys.argv[1:])
